import copy
from random import randint

import torch

from algorithms.algorithm import Algorithm
from algorithms.base.nsga2 import NSGA2
from individual.population import Population
from problems.problem import Problem
from utils.sort import sortrows
from utils.vectors import get_ref_vectors, unique


class MOCGDE(Algorithm):
    parameters = [
        {'label': 'Np', 'name': 'Np', 'type': 'number', 'description': 'Small population size', 'step': 1,
         'defaultValue': 10}
    ]

    def __init__(self, problem: Problem,
                 pop_size: int,
                 kwargs: dict,
                 ):
        super().__init__(pop_size, problem, kwargs)
        self.ref_vectors = get_ref_vectors(self.Np, problem.obj_dim)
        self.sub_n = self.ref_vectors.size(0)
        self.population = problem.init_population(self.sub_n)
        self.archive = copy.deepcopy(self.population)
        self.k = torch.zeros(self.sub_n)
        self.g0 = [None] * self.sub_n
        self.d0 = [None] * self.sub_n

    def final_output(self):
        return self.archive

    def each_iteration(self):
        self.k = torch.fmod(self.k, self.problem.var_dim) + 1
        off_pop = []
        dec = self.population.dec
        obj = self.population.obj
        archive_dec = self.archive.dec
        for i in range(self.sub_n):
            gk, site = self.finite_difference(self.population[i], self.problem, self.ref_vectors[i])
            if self.k[i] == 1:
                dk = - gk
            else:
                beta = (gk.T @ gk) / (self.g0[i].T @ self.g0[i])
                dk = -gk + beta * self.d0[i]
                if gk.T @ dk >= 0:
                    dk = -gk
            success = False
            for step in range(10):
                mu = torch.rand(1, self.problem.var_dim) < 1 / torch.sum(site)
                off_dec = dec[i] + ~site * (0.5 ** step) * dk.T + mu * site * (0.5 ** step) * (
                        archive_dec[randint(0, len(archive_dec) - 1)] - archive_dec[randint(0, len(archive_dec) - 1)])
                off_dec = self.problem.repair_decision(off_dec)
                solution = self.problem.estimate_population(off_dec)[0]
                off_obj = solution.obj
                off_pop.append(solution)
                if torch.all(off_obj < obj[i]):
                    success = True
                    break
            if success:
                self.population[i] = solution
                self.g0[i] = gk
                self.d0[i] = dk
            else:
                rand_index = randint(0, len(archive_dec) - 1)
                self.population[i] = self.archive[rand_index]
                self.k[i] = 0
        self.archive = self.update_archive(self.archive + Population(off_pop), self.pop_size)

    @staticmethod
    def finite_difference(solution, problem, ref_vectors):
        con = solution.con
        dec = solution.dec
        if con > 0:
            df = problem.calc_con_grad(dec)
            site = torch.zeros(1, len(dec), dtype=torch.bool)
            df = torch.sum(df, dim=1)
        else:
            df = problem.cal_grad(dec)
            site = torch.any(df < 0, dim=1) & torch.any(df > 0, dim=1)
            df = df @ ref_vectors.unsqueeze(dim=1)
        return df, site.unsqueeze(dim=0)

    @staticmethod
    def update_archive(population: Population, pop_size):
        dec = population.dec
        obj = population.obj
        index = unique(obj)
        dec = dec[index]
        obj = obj[index]
        front_no, _ = NSGA2.ND_sort(obj, 1)
        obj = obj[front_no == 1]
        dec = dec[front_no == 1]
        if len(obj) > pop_size:
            choose = torch.ones(len(obj), dtype=torch.bool)
            dis = torch.cdist(obj, obj)
            dis[torch.eye(len(dis), dtype=torch.bool)] = torch.inf
            while torch.sum(choose) > pop_size:
                remain = torch.nonzero(choose).squeeze(dim=1)
                temp = torch.sort(dis[remain], dim=1)[0]
                _, rank = sortrows(temp, return_sorted_indices=True)
                choose[remain[rank[0]]] = False
            obj = obj[choose]
            dec = dec[choose]
        return Population(dec, obj)
